import numpy as np
import pandas as pd
import rasterio
import matplotlib.pyplot as plt
from rasterio.warp import reproject, Resampling
from rasterio.enums import Resampling as RasterioResampling


def read_and_align_images(image1_path, image2_path):
    """读取并对齐两幅图像，使它们具有相同的形状和投影"""
    with rasterio.open(image1_path) as src1, rasterio.open(image2_path) as src2:
        # 检查两幅图像是否已经对齐
        if src1.shape == src2.shape: # and src1.transform == src2.transform:  # 两者维度和投影一样
            image1 = src1.read(1)
            image2 = src2.read(1)
            return image1, image2, src1.nodata, src2.nodata
        
        print(src1.shape)
        print(src2.shape)

        # 如果不匹配，将第二幅图像重采样到第一幅图像的网格上
        print("图像不匹配，正在重采样...")
        # 创建目标数组
        image1 = src1.read(1)
        image2_resampled = np.zeros_like(image1)
        
        # 重采样
        reproject(
            source=rasterio.band(src2, 1),
            destination=image2_resampled,
            src_transform=src2.transform,
            src_crs=src2.crs,
            dst_transform=src1.transform,
            dst_crs=src1.crs,
            resampling=Resampling.nearest
        )
        
        return image1, image2_resampled, src1.nodata, src2.nodata

def classify_and_sample(
    image1, image2, nodata1, nodata2, 
    num_bins=100, custom_bins=None, sample_frac=1.0, 
    output_csv=None, remove_outliers=True, remove_small_values=True,
    small_values= 0.1
):
    """
    根据image1的值域对image2进行采样并分类，保存每分级的均值和方差到CSV。
    :param num_bins: 分级数量
    :param custom_bins: 自定义分级边界
    :param sample_frac: 采样比例 (0-1)
    :param output_csv: 输出CSV文件的路径 (可选)
    :param remove_outliers: 是否去除异常值 (默认True)
    :param remove_small_values: 是否删除image2中值小于0.1的数据 (默认True)
    :param small_values: 默认，小于0.1的数据
    """
    
    # 创建有效值掩码
    mask1 = (image1 != nodata1) if nodata1 is not None else np.ones_like(image1, dtype=bool)
    mask2 = (image2 != nodata2) if nodata2 is not None else np.ones_like(image2, dtype=bool)
    valid_mask = mask1 & mask2
    
    # 获取有效值
    valid_image1 = image1[valid_mask]
    valid_image2 = image2[valid_mask]
    
    # 删除image2中值小于0.1的数据, small_values
    if remove_small_values:
        small_value_mask = valid_image2 >= small_values
        valid_image1 = valid_image1[small_value_mask]
        valid_image2 = valid_image2[small_value_mask]

    # 创建分级
    if custom_bins is not None:
        bins = custom_bins
    else:
        min_val = np.min(valid_image1)
        max_val = np.max(valid_image1)
        bins = np.linspace(min_val, max_val, num_bins+1)
    
    # 确定每个值属于哪个分级
    bin_indices = np.digitize(valid_image1, bins) - 1
    
    # 采样
    if sample_frac < 1.0:
        sample_mask = np.random.rand(len(valid_image2)) <= sample_frac
        sampled_image2 = valid_image2[sample_mask]
        sampled_bin_indices = bin_indices[sample_mask]
    else:
        sampled_image2 = valid_image2
        sampled_bin_indices = bin_indices
    
    # 按分级分组数据并计算统计量
    
    # 按分级分组数据
    grouped_data = []  # 分组原始数据
    bin_stats = []  # 分组统计数据
    bin_labels = []
    for i in range(len(bins)-1):
        mask = (sampled_bin_indices == i)
        if np.any(mask):
            values = sampled_image2[mask]
            grouped_data.append(sampled_image2[mask])

            # 去除异常值 (IQR方法)
            if remove_outliers and len(values) > 0:
                q1 = np.percentile(values, 25)
                q3 = np.percentile(values, 75)
                iqr = q3 - q1
                lower_bound = q1 - 1.5 * iqr
                upper_bound = q3 + 1.5 * iqr
                values = values[(values >= lower_bound) & (values <= upper_bound)]
            
            if len(values) > 0:
                mean = np.mean(values)
                variance = np.var(values)
                bin_labels.append(f"{bins[i]:.2f}-{bins[i+1]:.2f}")
                bin_stats.append({
                    "Bin": bin_labels[-1],
                    "Mean": mean,
                    "Variance": variance,
                    "Sample_Count": len(values)
                })
    
    # 保存为CSV文件
    if output_csv and bin_stats:
        df = pd.DataFrame(bin_stats)
        df.to_csv(output_csv, index=False, columns=["Bin", "Mean", "Variance", "Sample_Count"])
    
    
    return grouped_data, bin_labels

def plot_boxplot(grouped_data, bin_labels, title="Boxplot by Value Ranges", show_outliers=False):
    """绘制箱线图，支持调节异常值显示"""
    plt.figure(figsize=(12, 6))
    
    # 设置箱线图的异常值显示选项
    boxplot_props = {
        'showfliers': show_outliers,  # 是否显示异常值
        'showmeans': True,            # 显示平均值
        'meanline': True,            # 用线表示平均值
        'meanprops': {                # 平均值线的样式
            'color': 'red',
            'linestyle': '--',
            'linewidth': 2
        }
    }

    # 设置中文字体
    plt.rcParams['font.sans-serif'] = ['SimSun']
    plt.rcParams['axes.unicode_minus'] = False
    
    plt.boxplot(grouped_data, labels=bin_labels, **boxplot_props)
    plt.xticks(rotation=45, ha='right')
    plt.title(title)
    plt.tight_layout()
    plt.show()

def plot_scatter(image1, image2, nodata1, nodata2, 
                 sample_frac=1.0, title="Scatter Plot", 
                 xlabel="Image1 Values", ylabel="Image2 Values",
                 alpha=0.1, figsize=(10, 8)):
    """
    根据两张图像的有效值绘制散点图
    
    Parameters:
    -----------
    image1, image2: numpy arrays
        输入的两张图像数据
    nodata1, nodata2: float or None
        无效值标记
    sample_frac: float (0-1)
        采样比例
    title: str
        图表标题
    xlabel, ylabel: str
        x轴和y轴标签
    alpha: float
        点透明度
    figsize: tuple
        图表大小
    """
    # 创建有效值掩码
    mask1 = (image1 != nodata1) if nodata1 is not None else np.ones_like(image1, dtype=bool)
    mask2 = (image2 != nodata2) if nodata2 is not None else np.ones_like(image2, dtype=bool)
    # 新增：排除小于0.1的值
    mask1 = mask1 & (image1 >= 0.1)
    mask2 = mask2 & (image2 >= 0.1)
    valid_mask = mask1 & mask2
    
    # 获取有效值
    valid_image1 = image1[valid_mask]
    valid_image2 = image2[valid_mask]
    
    # 采样
    if sample_frac < 1.0:
        sample_size = int(len(valid_image1) * sample_frac)
        indices = np.random.choice(len(valid_image1), size=sample_size, replace=False)
        x = valid_image1[indices]
        y = valid_image2[indices]
    else:
        x = valid_image1
        y = valid_image2
    
    # 绘制散点图
    plt.figure(figsize=figsize)
    plt.scatter(x, y, alpha=alpha, s=5)  # s控制点的大小
    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.show()


# # 示例用法
# if __name__ == "__main__":
#     # # ----------------------------MODIS NDVI---------------------------------------
#     # 水热匹配指数WH与MODIS NDVI的关系
#     # 利用水热匹配指数WH=MAP/GDD(积温)，计算等值线
#     # file1 = r"D:\0GPPvsWater\Modis\atmp0_VS_pre.tif"
#     # file1 = r"D:\0GPPvsWater\ERA5-land\ERA5_pre_vs_atmp_2km.tif"
#     file1 = r"D:\0GPPvsWater\Modis\pre_VS_atmp0.tif"
#     file2 = r"D:\0GPPvsWater\Modis\Ndvi_MAX_2000_2024.tif"
#     small_values= 0.01
#     output_csv=".\WH_NDVI1.csv"

#     # 读取并对齐图像
#     image1, image2, nodata1, nodata2 = read_and_align_images(file1, file2)

#     # grouped_data, bin_labels = classify_and_sample(
#     #     image1, image2, nodata1, nodata2, 
#     #     custom_bins=custom_bins, sample_frac=0.5, output_csv=output_csv,
#     #     remove_outliers=True, remove_small_values=True, small_values=small_values
#     # )

#     # 选项1: 自动分成100份
#     grouped_data, bin_labels = classify_and_sample(
#         image1, image2, nodata1, nodata2, 
#         num_bins=100,sample_frac=0.5,
#         output_csv=output_csv, remove_outliers=True, remove_small_values=True, 
#         small_values=small_values)

#     # 绘制箱线图
#     plot_boxplot(grouped_data, bin_labels,title=file1)


# 示例用法
if __name__ == "__main__":
    # # 输入文件路径

    # # ----------------------------GIMMS NDVI---------------------------------------
    # file1 = r"D:\0GPPvsWater\ERA5-land\ERA5_pre_vs_atmp_2km.tif"
    # file2 = r"D:\0GPPvsWater\ERA5-land\ndvimean2_nodata.tif"
    # small_values= 0.1

    # # AI与NDVI的关系
    # file1 = r"D:\0GPPvsWater\散点图\AI.tif"
    # file2 = r"D:\0GPPvsWater\散点图\ndvimean_con.tif"
    # small_values= 0.1

    # # ----------------------------MODIS NDVI---------------------------------------
    # # MODIS 与NDVI的关系
    # file1 = r"D:\0GPPvsWater\Modis\atmp0_VS_pre.tif"
    # # file1 = r"D:\0GPPvsWater\Modis\pre_VS_atmp0.tif"
    # file2 = r"D:\0GPPvsWater\Modis\Ndvi_MAX_2000_2024.tif"
    # small_values= 0.1

    # # MODIS 与NDVI的关系 分大洲
    # file1 = r"D:\0GPPvsWater\Modis\分洲\atmp0_VS_pre\atmp0_VS_pre_South America.tif"
    # file2 = r"D:\0GPPvsWater\Modis\分洲\NDVI\NDVI_South America.tif"
    # small_values= 0.1

    # MODIS NDVI 每个像元的方差 与NDVI的关系
    file1 = r"D:\0GPPvsWater\Modis\atmp0_VS_pre.tif"
    # file1 = r"D:\0GPPvsWater\Modis\pre_VS_atmp0.tif"
    file2 = r"D:\0GPPvsWater\Modis\统计量\ndvi_variance_2000_2024.tif"
    small_values= 0.000001
    
    # 读取并对齐图像
    image1, image2, nodata1, nodata2 = read_and_align_images(file1, file2)
    
    # # 选项1: 自动分成100份
    # grouped_data, bin_labels = classify_and_sample(
    #     image1, image2, nodata1, nodata2, 
    #     num_bins=100, sample_frac=0.1  # 采样10%的数据
    # )
    
    # 选项2: 自定义分级范围
    # custom_bins = [0,0.1,0.2, 0.3,0.4, 0.5,0.6,0.7,0.8,0.9,1.0,1.1,1.2,1.3,1.4,1.5,1.6,1.7,1.8,1.9,
    #                2.0,2.1,2.2,2.3,2.4,2.5,2.6,2.7,2.8,2.9,3.0,3.1,3.2,3.3,3.4,3.5,3.6,3.7,3.8,3.9,
    #                4.0,4.1,4.2,4.3,4.4,4.5,4.6,4.7,4.8,4.9,5.0,5.1,5.2,5.3,5.4,5.5,5.6,5.7,5.8,5.9,
    #                6.0,6.1,6.2,6.3,6.4,6.5,6.6,6.7,6.8,6.9,7.0,8.0,9.0, 10, 20, 50, 100,1000,10000,100000]  # 示例自定义分级
    

    # 旧分区 3.95
    # custom_bins = [0,0.1975,0.395,0.5925,0.79,0.9875,1.185,1.3825,1.58,1.7775,1.975,2.1725,2.37,2.5675,2.765,
    #                2.9625,3.16,3.3575,3.555,3.7525,3.95,4.157894737,4.388888889,4.647058824,4.9375,5.266666667,
    #                5.642857143,6.076923077,6.583333333,7.181818182,7.9,8.777777778,9.875,11.28571429,13.16666667,
    #                15.8,19.75,26.33333333,39.5,79,100000]

    # 新分区 4.94
    custom_bins = [ 0.247, 0.494, 0.741, 0.988, 1.235, 1.482, 1.729, 1.976, 2.223, 2.47, 2.717, 2.964, 3.211, 3.458, 
                   3.705, 3.952, 4.199, 4.446, 4.693, 4.94, 5.200, 5.489, 5.812, 6.175, 6.587, 7.057, 7.600, 8.233, 
                   8.982, 9.880, 10.978, 12.350, 14.114, 16.467, 19.760, 24.700, 32.933, 49.400, 98.800, 100000 ]
    
    grouped_data, bin_labels = classify_and_sample(
        image1, image2, nodata1, nodata2, 
        custom_bins=custom_bins, sample_frac=0.5, output_csv=".\水热\output_data.csv",
        remove_outliers=True, remove_small_values=True, small_values=small_values
    )
    
    # 绘制箱线图
    plot_boxplot(grouped_data, bin_labels,title=file1)

        # 绘制散点图
    # plot_scatter(image1, image2, nodata1, nodata2,
    #              sample_frac=0.1,  # 采样10%的数据
    #              title="Scatter Plot of Image1 vs Image2",
    #              xlabel="Image1 Values",
    #              ylabel="Image2 Values")
